from torchvision.datasets import FashionMNIST
from torchvision import transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "True"

# 下载数据集
train_data = FashionMNIST(root='./data',
                          train=True,
                          download=True,
                          transform=transforms.Compose([transforms.Resize(size=224), transforms.ToTensor()]))

# 加载数据集
train_loader = DataLoader(train_data,
                          batch_size=64,
                          shuffle=True)

# 获得一个Batch的数据
for step, (x, y) in enumerate(train_loader):
    if step > 0:
        break
batch_x = x.squeeze().numpy()  # 将四维张量移除第一维，转为numpy数组 [B W H C]
batch_y = y.numpy()
label = train_data.classes

# 可视化一个Batch的图像
plt.figure(figsize=(18, 6))
for i in range(len(batch_y)):
    plt.subplot(4, 16, i+1)
    plt.imshow(batch_x[i, :, :])
    plt.title(label[batch_y[i]],size=10)
    plt.axis("off")
    plt.subplots_adjust(wspace=0.05)
plt.show()